import base64
from typing import TypedDict
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
from cryptography.hazmat.primitives import serialization


class ECDHKeyPair(TypedDict):
    publicKey: ec.EllipticCurvePrivateKey
    privateKey: ec.EllipticCurvePrivateKey
    publicKeyBytes: bytes
    publicKeyBase64: str
    privateKeyBytes: bytes
    privateKeyBase64: str


MYECDH_ENCRYPTION_INFO = "ECDH key derivation"
MYECDH_ENCRYPTION_SALT = "1234"


def generate_key_pair() -> ECDHKeyPair:
    # 生成 ECDH 密钥对 (使用 P-256 曲线)
    private_key: ec.EllipticCurvePrivateKey = ec.generate_private_key(
        ec.SECP256R1(), default_backend()
    )
    public_key: ec.EllipticCurvePublicKey = private_key.public_key()
    # 序列化公钥为 uncompressed 格式
    public_key_bytes: bytes = public_key.public_bytes(
        encoding=serialization.Encoding.X962,
        format=serialization.PublicFormat.UncompressedPoint,
    )
    # 序列化私钥为 PKCS8 格式
    private_key_bytes: bytes = private_key.private_bytes(
        encoding=serialization.Encoding.DER,
        format=serialization.PrivateFormat.PKCS8,
        encryption_algorithm=serialization.NoEncryption(),
    )
    return {
        "privateKey": private_key,
        "publicKey": public_key,
        "publicKeyBytes": public_key_bytes,
        "publicKeyBase64": base64.b64encode(public_key_bytes).decode("utf-8"),
        "privateKeyBytes": private_key_bytes,
        "privateKeyBase64": base64.b64encode(private_key_bytes).decode("utf-8"),
    }


def compute_shared_secret(
    private_key: ec.EllipticCurvePrivateKey, peer_public_key_bytes: bytes
) -> bytes:
    peer_public_key: ec.EllipticCurvePublicKey = (
        ec.EllipticCurvePublicKey.from_encoded_point(
            ec.SECP256R1(), peer_public_key_bytes
        )
    )
    shared_secret: bytes = private_key.exchange(ec.ECDH(), peer_public_key)
    return shared_secret


def compute_shared_secret_via64(
    private_key_base64: str, peer_public_key_base64: str
) -> str:
    private_key_bytes: bytes = base64.b64decode(private_key_base64)
    private_key = serialization.load_der_private_key(
        private_key_bytes, password=None, backend=default_backend()
    )
    peer_public_key_bytes = base64.b64decode(peer_public_key_base64)
    shared_secret: bytes = compute_shared_secret(private_key, peer_public_key_bytes)
    return base64.b64encode(shared_secret).decode("utf-8")


def compute_derive_key(
    shared_secret: bytes,
    salt: str = MYECDH_ENCRYPTION_SALT,
    info: str = MYECDH_ENCRYPTION_INFO,
) -> bytes:
    hkdf = HKDF(
        algorithm=hashes.SHA256(),
        length=32,
        salt=salt.encode("utf-8"),
        info=info.encode("utf-8"),
        backend=default_backend(),
    )
    derived_key = hkdf.derive(shared_secret)
    return derived_key


def compute_derive_key_viab64(
    shared_secret_base64: str,
    salt: str = MYECDH_ENCRYPTION_SALT,
    info: str = MYECDH_ENCRYPTION_INFO,
) -> str:
    shared_secret: bytes = base64.b64decode(shared_secret_base64)
    derived_key: bytes = compute_derive_key(shared_secret, salt, info)
    return base64.b64encode(derived_key).decode("utf-8")


def test():
    print("ECDH 密钥交换演示 (服务端)")
    print("生成密钥对...")
    # 生成密钥对
    kp: ECDHKeyPair = generate_key_pair()

    print("\n公钥 (Base64):")
    print(kp["publicKeyBase64"])
    print("\n私钥 (Base64):")
    print(kp["privateKeyBase64"])

    # 等待用户输入客户端公钥
    print("\n请输入客户端的公钥 (Base64):")
    input_public_key_b64 = input().strip()

    # 解码客户端公钥
    # client_public_key_bytes = base64.b64decode(input_public_key_b64)
    # 计算共享密钥
    # shared_secret: bytes = compute_shared_secret(
    #    kp["privateKey"], client_public_key_bytes
    # )
    shared_secret = compute_shared_secret_via64(
        kp["privateKeyBase64"], input_public_key_b64
    )

    print("\n共享密钥 (Base64):")
    # print(base64.b64encode(shared_secret).decode("utf-8"))
    print(shared_secret)
    # 派生密钥

    print("\n派生密钥 (Base64):")
    print(compute_derive_key_viab64(shared_secret, "1234", "ECDH key derivation"))
    ######################################
    shared_secret: bytes = base64.b64decode(shared_secret)
    derived_key: bytes = compute_derive_key(shared_secret)
    print(base64.b64encode(derived_key).decode("utf-8"))


if __name__ == "__main__":
    test()
